#!/usr/bin/python
# -*- coding:utf-8 -*-

"""任务数据迁移"""

import datetime
from app.db.models.jobs import Job, VulSettings, SslSettings, ContentSettings, SecurityEventSettings
from app.db.models.tasks import Task
from app.db.models.asset_tasks import AssetTasks
from app.db.models.assets import Asset


class RewriteAssetTask(object):
    """
    任务数据迁移， 把jobs表中的周期任务配置同步添加到asset_tasks表
    把tasks表中的单次任务数据同步到asset_tasks表

    每个job 的vulSettings, sslSettings, contentSettings, securityEventSettings 组合成一个web安全监测周期任务
    根据每个job 的pingSettings 生成一个或多个ping周期任务
    根据每个job 的httpSettings 生成一个或多个http周期任务
    每个asset 数据生成一个资产变动任务
    tasks 表中的单次ping、http 任务忽略， 不处理
    每个tasks 表中的 vul, ssl, content, securityEvent 单次任务 分别生成一个对应的web安全监测单次任务
    """

    def __init__(self, last_task_id: str = None, job_ids: list = None):
        """
        last_task_id 只会同步id > last_task_id 的单次任务
        job_ids 只会同步指定的资产id 数据
        """
        self.job_ids = job_ids
        self.last_task_id = last_task_id
        self.jobs_info = {}
        self.web_sec_schedule_task = []
        self.http_schedule_task = []
        self.ping_schedule_task = []
        self.asset_schedule_task = []
        self.web_sec_manual_task = []
        self.asset_task_ids = []
        self.exists_schedule_task_job_ids = self._exists_schedule_task_job()
        self.exists_manual_task_ids = self._exists_manual_task()

    @staticmethod
    def _exists_schedule_task_job():
        print(f"{datetime.datetime.now()} 统计已经迁移过周期任务的资产......")
        return {a.jobId: 1 for a in AssetTasks.objects.filter(triggerType="schedule")}

    @staticmethod
    def _exists_manual_task():
        print(f"{datetime.datetime.now()} 统计已经迁移web安全监测单次任务......")
        task_ids = []
        for at in AssetTasks.objects.filter(triggerType="manual", taskGroupType="web_security"):
            for t in ["vul", "ssl", "securityEvent", "content"]:
                if task_id := getattr(at, f"{t}Settings").taskId:
                    task_ids.append(task_id)
        return [task.id for task in Task.objects.filter(taskId__in=task_ids).only("id")]

    def parse_job(self):
        query = {"assetType__ne": "host",
                 "verification__isVerified": True}
        if self.job_ids:
            query.update({"id__in": self.job_ids})
        task_common_fields = ["Settings", "AlertState", "ResultId", "SecurityStatus"]
        avai_task_fields = ["PrevResultId", "Abnormal", "NextExecDate", *task_common_fields]
        web_sec_fields = [*task_common_fields, "IsStarted"]
        for job in Job.objects.filter(**query):
            self.jobs_info.update({job.id: {"targetUrl": job.targetUrl,
                                            "note": job.note,
                                            "enableMonitor": job.enableMonitor,
                                            "alertSettings": job.alertSettings
                                            }})
            if self.exists_schedule_task_job_ids.get(job.id):
                print(f"{job.id}: {job.targetUrl}已经同步过周期任务")
                continue
            common_fields = {
                "targetUrl": job.targetUrl,
                "uid": job.uid,
                "jobId": job.id,
                "isVerified": job.verification.isVerified,
                "triggerType": "schedule",
                "enableMonitor": job.enableMonitor,
                "taskType": "",
                "alertSettings": job.alertSettings,
                "customMonitorPage": job.customMonitorPage,
                "prelogin": job.prelogin
            }
            web_sec_task = dict(contentPrevResultId=job.contentPrevResultId,
                                sitePortraitId=job.sitePortraitId,
                                taskGroupType="web_security",
                                name=job.note or "未命名",
                                **common_fields)
            web_task_ids = []
            web_enable_list = []
            for web_sec_type in ["vul", "ssl", "securityEvent", "content"]:
                for key in web_sec_fields:
                    attr = f"{web_sec_type}{key}"
                    if obj := getattr(job, attr, ""):
                        web_sec_task.update({attr: obj})
                    if key == "Settings" and obj:
                        web_task_ids.append(obj.taskId)
                        web_enable_list.append(obj.enable)
            if not any(web_enable_list):
                web_sec_task.update({"enableMonitor": False})
            if any(web_task_ids):
                self.web_sec_schedule_task.append(web_sec_task)
            self.parse_schedule_avai_task("ping", job, avai_task_fields, common_fields)
            self.parse_schedule_avai_task("http", job, avai_task_fields, common_fields)

    def parse_asset(self):
        assets = Asset.objects.filter(jobId__in=list(self.jobs_info.keys()))
        for asset in assets:
            if self.exists_schedule_task_job_ids.get(asset.jobId):
                continue
            asset_task = {
                "targetUrl": self.jobs_info[asset.jobId]["targetUrl"],
                "name": self.jobs_info[asset.jobId]["note"] or "未命名",
                "uid": asset.uid,
                "jobId": asset.jobId,
                "isVerified": True,
                "triggerType": asset.triggerType,
                "enableMonitor": True,
                "taskType": "",
                "alertSettings": self.jobs_info[asset.jobId]["alertSettings"],
                "taskGroupType": "web_asset",
                "assetSecurityStatus": asset.securityStatus,
                "assetAlertState": asset.alertState,
                "assetResultId": asset.resultId,
                "assetSettings": asset.taskSettings
            }
            self.asset_schedule_task.append(asset_task)

    def parse_schedule_avai_task(self, task_type, job, avai_task_fields, common_fields):
        avai_dict = {}
        for key in avai_task_fields:
            attr = f"{task_type}{key}"
            if obj := getattr(job, attr, ""):
                avai_dict.update({attr: obj})
        if length := len(avai_dict.get(f"{task_type}Settings", [])):
            for i in range(length):
                d = {"taskGroupType": task_type}
                task_id = ""
                for key, settings in avai_dict.items():
                    try:
                        d.update({key: settings[i]})
                        if key.endswith("Settings"):
                            task_id = settings[i].taskId
                            enable = settings[i].enable
                            d.update({"name": settings[i].name, "enableMonitor": enable})
                        if key.endswith("AlertState"):
                            d.update({key: settings[i].to_dict()})
                    except IndexError:
                        pass
                if task_id:
                    d.update(common_fields)
                    getattr(self, f"{task_type}_schedule_task").append(d)

    def add_schedule(self):
        tasks = self.http_schedule_task + self.web_sec_schedule_task + self.ping_schedule_task + self.asset_schedule_task
        for task in tasks:
            a = AssetTasks(**task)
            a.save()
            if a.taskGroupType not in ["http", "ping"]:
                # 更新运行状态
                a.update_status()
                # 更新任务结果数据
                a.get_web_security_last_result()
            self.asset_task_ids.append(a.id)

    def add_manual(self):
        only_fields = ["id", "name", "uid", "jobId", "alertSettings", "taskType",
                       "prelogin", "taskSettings", "securityStatus", "status", "taskId"]
        query = {
            "jobId__in": list(self.jobs_info.keys()),
            "taskType__in": ["vul", "ssl", "securityEvent", "content"],
            "triggerType": "manual",
            "id__nin": self.exists_manual_task_ids
        }
        if self.last_task_id:
            query.update({"id__gt": self.last_task_id})
        tasks = Task.objects.filter(**query).only(*only_fields)
        for task in tasks:
            web_sec_task = {
                "targetUrl": self.jobs_info[task.jobId]["targetUrl"],
                "name": task.name or "未命名",
                "uid": task.uid,
                "jobId": task.jobId,
                "isVerified": True,
                "triggerType": "manual",
                "enableMonitor": True,
                "taskType": "",
                "alertSettings": task.alertSettings,
                "prelogin": task.prelogin,
                "taskGroupType": "web_security",
                f"{task.taskType}SecurityStatus": task.securityStatus,
            }
            if task.name == "源站安全检测":
                web_sec_task.update({"is_defense_task": True})
            if task.status == "completed":
                web_sec_task.update({f"{task.taskType}ResultId": task.id})
            task_settings_dict = task.taskSettings.to_dict()
            task_settings_dict.update({"taskId": task.taskId, "enable": True})
            if task.taskType == "vul":
                web_sec_task.update({"vulSettings": VulSettings.from_dict(task_settings_dict)})
            elif task.taskType == "ssl":
                web_sec_task.update({"sslSettings": SslSettings.from_dict(task_settings_dict)})
            elif task.taskType == "content":
                web_sec_task.update({"contentSettings": ContentSettings.from_dict(task_settings_dict)})
            elif task.taskType == "securityEvent":
                web_sec_task.update({"securityEventSettings": SecurityEventSettings.from_dict(task_settings_dict)})
            else:
                continue
            a = AssetTasks(**web_sec_task)
            a.save()
            # 更新运行状态
            a.update_status()
            # 更新任务结果数据
            a.get_web_security_last_result()
            self.asset_task_ids.append(a.id)

    def run(self):
        print(f"{datetime.datetime.now()} 开始同步数据......")
        print(f"{datetime.datetime.now()} 解析web安全监测、可用性监测任务数据......")
        self.parse_job()
        print(f"{datetime.datetime.now()} 将增加web安全周期任务{len(self.web_sec_schedule_task)}个!")
        print(f"{datetime.datetime.now()} 将增加ping周期任务{len(self.ping_schedule_task)}个!")
        print(f"{datetime.datetime.now()} 将增加http周期任务{len(self.http_schedule_task)}个!")
        print(f"{datetime.datetime.now()} 解析资产变动任务数据......")
        self.parse_asset()
        print(f"{datetime.datetime.now()} 将增加资产变动周期任务{len(self.asset_schedule_task)}个!")
        try:
            print(f"{datetime.datetime.now()} 写入周期任务数据......")
            self.add_schedule()
            schedule_count = len(self.asset_task_ids)
            print(f"{datetime.datetime.now()} 写入周期任务{schedule_count}个!")
            print(f"{datetime.datetime.now()} 写入单次任务数据......")
            self.add_manual()
            print(f"{datetime.datetime.now()} 写入单次任务{len(self.asset_task_ids) - schedule_count}个!")
        except:
            import traceback
            traceback.print_exc()
            print(f"{datetime.datetime.now()} 同步失败, 回退数据......")
            AssetTasks.objects(id__in=self.asset_task_ids).delete()
            print(f"{datetime.datetime.now()} 回退完成！！！")
        else:
            print(f"{datetime.datetime.now()} 同步完成, 共同步{len(self.asset_task_ids)}条数据到asset_tasks")
        print(f"{datetime.datetime.now()} 结束！！！")


if __name__ == "__main__":
    RewriteAssetTask(job_ids=["632d7776dbec40c6aa7dba5a", ]).run()
